# Copyright (c) 2020 Presto Labs Pte. Ltd.
# Author: jhkim

import os
import numpy
from absl import app, flags

from coin.strategy.mm.tool.feed_stat_util import enumerate_volume_trend
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
from matplotlib.backends.backend_pdf import PdfPages
from pandas.plotting import register_matplotlib_converters
register_matplotlib_converters()


def main(_):
  trading_date_range = flags.FLAGS.trading_date.split("-")
  trading_date_from = int(trading_date_range[0])
  trading_date_to = int(trading_date_range[-1])
  os.makedirs(flags.FLAGS.out_dir, exist_ok=True)
  for nametuple, df in enumerate_volume_trend(trading_date_from, trading_date_to, symbolgroup=True):
    plt.rcParams['figure.figsize'] = 8, 4
    plt.rcParams['font.size'] = 8
    plt.rcParams['legend.fontsize'] = 8
    plt.rcParams['xtick.labelsize'] = 8
    plt.rcParams['ytick.labelsize'] = 8
    plt.rcParams['lines.linewidth'] = 0.5

    idx_ordered = df.groupby(["exchange", "market_type",
                              "symbolgroup"]).sum().sort_values(['total_notional'],
                                                                ascending=False).index
    groupdict = {
        symbolgroup: subdf for symbolgroup,
        subdf in df.groupby(["exchange", "market_type", "symbolgroup"])
    }

    with PdfPages(f"{flags.FLAGS.out_dir}/{nametuple}.pdf") as pdf:
      for symbolgroup in idx_ordered:
        subdf = groupdict[symbolgroup]
        plt.plot(subdf['trading_date'],
                 subdf['total_notional'],
                 'r.-',
                 subdf['trading_date'],
                 subdf['total_notional'].rolling(10).mean(),
                 'b:',
                 numpy.nan,
                 numpy.nan,
                 'g:',
                 drawstyle='steps-post',
                 markersize=2,
                 linewidth=0.5)
        plt.legend(['volume', 'volume 10d ma', 'volume 10d ma logscale'])
        plt.title(symbolgroup)
        plt.xticks(rotation=20)
        plt.twinx()
        plt.plot(subdf['trading_date'], numpy.log(subdf['total_notional'].rolling(10).mean()), 'g:')
        pdf.savefig()
        plt.close()


if __name__ == '__main__':
  flags.DEFINE_string('out_dir', 'volume_trend', '')
  flags.DEFINE_string('trading_date', None, '')
  app.run(main)
